import os
from openpyxl import load_workbook, Workbook
import json
import pprint as pp
import shutil
import re


def get_ext(file_path):
    """
    获取文件的扩展名

    :param file_path: 文件路径
    :return: 文件的扩展名，包括点号（.）
    """
    # 使用os.path.splitext函数将文件路径分割成文件名和扩展名两部分
    # 返回的结果是一个元组，第二个元素就是扩展名
    return os.path.splitext(file_path)[1]


def get_files(dir: str, exts=[], recursive=False) -> []:
    """
    获取指定目录下指定扩展名的文件列表，默认不递归
    :param dir: 目录路径
    :param exts: 过滤扩展名，只列出数组中的文件，如果不设置，为所有文件。要带点 `.xlsx`
    :param recursive: 是否递归查找
    :return: 文件列表
    """
    file_paths = []

    if recursive:
        # 遍历指定目录及其子目录下的所有文件
        for root, _, files in os.walk(dir):
            for file in files:
                # 使用os.path.join拼接完整的文件路径
                file_paths.append(os.path.join(root, file))
    else:
        # 只取一层文件夹的文件
        file_paths = map(lambda x: os.path.join(dir, x), os.listdir(dir))

    def filter_fn(file_path):
        if len(exts) == 0:
            return True
        return get_ext(file_path) in exts

    # 排除扩展名不匹配的文件，还有 `~$` 开头的临时文件
    valid_ext_files = filter(
        lambda x: filter_fn(x) and not os.path.basename(x).startswith("~$"),
        file_paths,
    )

    return list(valid_ext_files)


def divide_array(arr, step):
    """将数组分割成指定步长的多个数组"""
    return [arr[i : i + step] for i in range(0, len(arr), step)]


def clean_text(text):
    return re.sub(r"\s+", "", str(text)) if text is not None else ""


def read_from_xlsx(src_file, sheet_index=0, row_offset=0, max_col=0, data_only=True):
    """从 xlsx 文件读取数据，全是0索引"""
    wb = load_workbook(src_file, data_only=data_only)
    ws = wb.worksheets[sheet_index]
    data = list(ws.values)[row_offset:]
    res = []
    for row in data:
        sub_row = row if max_col == 0 else row[:max_col]
        res.append(sub_row)
    return res


def write_to_xlsx(file_path, data):
    wb = Workbook()

    # 第一个工作表
    ws = wb.active
    for row in data:
        ws.append(row)

    for row in ws.iter_rows(min_row=1, max_col=100, max_row=200):
        for cell in row:
            cell.number_format = "@"

    wb.save(file_path)


def read_from_wb(wb, sheet_index=0, row_offset=0, max_col=0):
    """从 xlsx 文件读取数据，全是0索引"""
    ws = wb.worksheets[sheet_index]
    data = list(ws.values)[row_offset:]
    res = []
    for row in data:
        sub_row = row if max_col == 0 else row[:max_col]
        res.append(sub_row)
    return res


def build_city_teams(code="00"):
    """生成省份、城市队伍"""
    dir = os.path.dirname(__file__)
    area_file = os.path.join(dir, "area.json")
    with open(area_file, "r", encoding="utf-8") as f:
        data = json.load(f)
        pp.pprint(data["00"])
        citys = data[code]
        return list(citys.values())


def build_class_teams(count):
    """生成固定数量的班级名字"""
    return list(map(lambda x: f"{x}班", list(range(1, count + 1))))


def try_file_path(file_path):
    """测试是否存在文件，如果存在，返回带括号的文件名"""
    tmp_path = file_path
    (root, ext) = os.path.splitext(file_path)
    i = 1
    while os.path.exists(tmp_path):
        tmp_path = f"{root}({i}){ext}"
        i += 1

    return tmp_path


def prepare_dir(dir: str, renew=True) -> None:
    """
    准备目标文件夹：
    - 存在的话
        - 如果要更新，先删除再新建
        - 如果不更新，不需要新建
    - 不存在的话，新建
    :param dir: 目标文件夹路径
    :param renew: 是否删除重新创建
    """
    if os.path.exists(dir):
        if renew:
            shutil.rmtree(dir)
            os.mkdir(dir)
    else:
        os.mkdir(dir)
